import numpy
from linecache import getline

# 高程模型文件
source = "Grid/dem/dem.asc"
# 坡度网格文件
slopegrid = "Grid/dem/slope.asc"
# 坡向网格文件
aspectgrid = "Grid/dem/aspect.asc"
# 输出渲染的地图文件
shadegrid = "Grid/dem/shade.asc"

# 使用内置的linecache模块循环解析头部文件
hdr = [getline(source, i) for i in range(1, 7)]
values = [float(h.split()[-1].strip()) for h in hdr]
cols, rows, lx, ly, cell, nd = values
xres = cell
yrex = cell * -1

# 将dem载入numpy数组中
arr = numpy.loadtxt(source, skiprows=6)

# 排除围绕边框之外的2像素NODATA数据
# 同时设立3*3的窗口进行坡度计算
window = []
for row in range(3):
    for col in range(3):
        window.append(arr[row:(row + arr.shape[0] - 2), col:(col + arr.shape[1] - 2)])

# 处理水平和竖直方向上每个3 * 3窗体
# 高程放大比例
z = 1.0
# 分辨率
scale = 1.0
x = ((z*window[0] + z*window[3] + z*window[3] + z*window[6]) -
     (z*window[2] + z*window[5] + z*window[5] + z*window[8]))/(8.0 * xres * scale)
y = ((z*window[6] + z*window[7] + z*window[7] + z*window[8]) -
     (z*window[0] + z*window[1] + z*window[1] + z*window[2]))/(8.0 * yrex * scale)

rad2deg = 180.0/3.141592653589793
deg2rad = 3.141592653589793/180.0
# 计算坡度
slope = 90.0 - numpy.arctan(numpy.sqrt(x * x + y * y)) * rad2deg
# 计算坡度
aspect = numpy.arctan2(x, y)

# 计算晕渲阴影
# 日光角度
altitude = 45.0
# 日光方位
azimuth = 315.0
shaded = numpy.sin(altitude*deg2rad) * numpy.sin(slope*deg2rad) + \
    numpy.cos(altitude*deg2rad) * numpy.cos(slope * deg2rad) * \
    numpy.cos((azimuth - 90.0) * deg2rad - aspect)
# 阴影放大比例，可选值区间为： 0~1 到 0~255
shaded = shaded * 255
# 生成新的头部信息
header = "ncols        {}\n".format(shaded.shape[1])
header += "nrows        {}\n".format(shaded.shape[0])
header += "xllcorner        {}\n".format(lx + (cell * (cols - shaded.shape[1])))
header += "yllcorner        {}\n".format(ly + (cell * (rows - shaded.shape[0])))
header += "cellsize        {}\n".format(cell)
header += "NODATA        {}\n".format(-9999)  # -9999 为占位之，无意义

# 为窗体设定NODATA值
for pane in window:
    slope[pane == nd] = -9999
    aspect[pane == nd] = -9999
    shaded[pane == nd] = -9999

# 打开输出文件，添加头部信息，保存坡度网格
with open(slopegrid, "wb") as f:
    f.write(bytes(header, "UTF-8"))
    numpy.savetxt(f, slope, fmt="%4i")
# 打开输出文件，添加头部信息，保存坡向网格
with open(aspectgrid, "wb") as f:
    f.write(bytes(header, "UTF-8"))
    numpy.savetxt(f, aspect, fmt="%4i")
# 打开输出文件，添加头部信息，保存晕渲网格
with open(shadegrid, "wb") as f:
    f.write(bytes(header, "UTF-8"))
    numpy.savetxt(f, shaded, fmt="%4i")

